﻿/**
封装MP4
*/

#include <iostream>

extern "C"
{
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavcodec/bsf.h>
#include <libswscale/swscale.h>
#include <libavutil/common.h>
#include <libavutil/imgutils.h>
#include <libavutil/timestamp.h>
};

#include <magic_enum_all.hpp>

//#define REMOVE_NALU

int main()
{
	std::string rtsp_url = "rtsp://admin:cdtye123456@192.168.116.201:554/Streaming/Channels/103";
	std::string out_file_path = "test3.mp4";

	AVPacket* av_packet = av_packet_alloc();
	AVPacket* bsf_pkt = av_packet_alloc();
	AVFormatContext* in_fmt_ctx = NULL;
	const AVOutputFormat* out_fmt = NULL;
	AVFormatContext* out_fmt_ctx = NULL;
	int ret = 0;

#ifdef REMOVE_NALU
	// 1 av_bsf_get_by_name
	const AVBitStreamFilter* bsfilter = av_bsf_get_by_name("h264_mp4toannexb");
	AVBSFContext* bsf_ctx = NULL;
	// 2 初始化过滤器上下⽂
	av_bsf_alloc(bsfilter, &bsf_ctx); //AVBSFContext;
#endif


	AVDictionary* options = NULL;
	av_dict_set(&options, "buffer_size", "4096000", 0); //设置缓存大小
	av_dict_set(&options, "rtsp_transport", "tcp", 0);  //以tcp的方式打开,
	av_dict_set(&options, "stimeout", "5000000", 0);    //设置超时断开链接时间，单位us,   5s
	av_dict_set(&options, "timeout", "5000000", 0);    //设置超时断开链接时间，单位us,   5s
	av_dict_set(&options, "max_delay", "500000", 0);    //设置最大时延

	//打开网络流或文件流
	if (avformat_open_input(&in_fmt_ctx, rtsp_url.c_str(), NULL, &options) != 0)
	{
		std::cout << "Couldn't open input stream.\n";
		goto finish;
	}

	//获取视频文件信息
	if (avformat_find_stream_info(in_fmt_ctx, NULL) < 0)
	{
		std::cout << "Couldn't find stream information." << std::endl;
		goto finish;
	}

	//打印输入流的信息
	av_dump_format(in_fmt_ctx, 0, rtsp_url.c_str(), 0);

	//查找码流中是否有视频流
	int video_index = -1;
	for (unsigned i = 0; i < in_fmt_ctx->nb_streams; i++)
	{
		if (in_fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			video_index = i;
			break;
		}
	}
	if (video_index == -1)
	{
		std::cerr << "Didn't find a video stream.\n";
		goto finish;
	}

	//输出context
	avformat_alloc_output_context2(&out_fmt_ctx, NULL, NULL, out_file_path.c_str());
	if (!out_fmt_ctx)
	{
		std::cout << "Couldn't create output context.\n";
		goto finish;
	}

	out_fmt = out_fmt_ctx->oformat;

	AVStream* out_stream = NULL;
	AVStream* in_stream = in_fmt_ctx->streams[video_index];
	AVCodecParameters* in_codecpar = in_stream->codecpar;

	//添加流
	out_stream = avformat_new_stream(out_fmt_ctx, NULL);
	if (!out_stream)
	{
		std::cerr << "failed allocating output stream.\n";
		goto finish;
	}
	//拷贝参数
	ret = avcodec_parameters_copy(out_stream->codecpar, in_codecpar);
	if (ret < 0)
	{
		std::cerr << "failed to copy codec parameters\n";
		goto finish;
	}

#ifdef REMOVE_NALU
	// 3 添加解码器属性
	avcodec_parameters_copy(bsf_ctx->par_in, in_fmt_ctx->streams[video_index]->codecpar);
	av_bsf_init(bsf_ctx);
#endif

	//自动选择
	out_stream->codecpar->codec_tag = 0;
	out_stream->r_frame_rate = in_stream->r_frame_rate;
	av_dump_format(out_fmt_ctx, 0, out_file_path.c_str(), 1);

	//文件，打开文件
	if (!(out_fmt->flags & AVFMT_NOFILE))
	{
		ret = avio_open(&out_fmt_ctx->pb, out_file_path.c_str(), AVIO_FLAG_WRITE);
		if (ret < 0)
		{
			std::cerr << "Could not open output file\n";
			goto finish;
		}
	}

	//写入文件头
	ret = avformat_write_header(out_fmt_ctx, NULL);
	if (ret < 0)
	{
		fprintf(stderr, "Error occurred when opening output file\n");
		goto finish;
	}

	int frame_index = 0;

	//开始循环读取数据
	while (frame_index < 250)
	{
		AVStream* in_stream = nullptr;
		AVStream* out_stream = nullptr;
		if (av_read_frame(in_fmt_ctx, av_packet) >= 0)
		{
			if (av_packet->stream_index == video_index)
			{
				out_stream = out_fmt_ctx->streams[0];
				in_stream = in_fmt_ctx->streams[video_index];
				av_packet_rescale_ts(av_packet, in_stream->time_base, out_stream->time_base);
				av_packet->pos = -1;//?

				//if (av_packet->pts == AV_NOPTS_VALUE )
				{
					AVRational time_base = out_stream->time_base;
					int64_t duration = (double)AV_TIME_BASE / (av_q2d(out_stream->r_frame_rate));

					av_packet->pts = (double)(frame_index * duration) / (double)(av_q2d(time_base) * AV_TIME_BASE);
					av_packet->dts = av_packet->pts;
					av_packet->duration = duration / (double)(av_q2d(time_base) * AV_TIME_BASE);
					av_packet->stream_index = out_stream->index;
				}

#ifdef REMOVE_NALU
				if (av_bsf_send_packet(bsf_ctx, av_packet) == 0) // bitstreamfilter内部去维护内存空间
				{
					if (av_bsf_receive_packet(bsf_ctx, bsf_pkt) == 0)
					{
						ret = av_interleaved_write_frame(out_fmt_ctx, bsf_pkt);

						if (ret < 0)
						{
							fprintf(stderr, "Error muxing packet\n");
							break;
						}
					}
				}
#else
				ret = av_interleaved_write_frame(out_fmt_ctx, av_packet);

				if (ret < 0)
				{
					fprintf(stderr, "Error muxing packet\n");
					break;
				}
#endif

				frame_index++;
				std::cout << frame_index << std::endl;
			}
		}
		av_packet_unref(av_packet);
	}

	av_write_trailer(out_fmt_ctx);

finish:
	if (options)
		av_dict_free(&options);
	av_packet_free(&av_packet);
	if (in_fmt_ctx)
		avformat_close_input(&in_fmt_ctx);

	if (out_fmt_ctx && !(out_fmt_ctx->flags & AVFMT_NOFILE))
	{
		avio_closep(&out_fmt_ctx->pb);
	}
	avformat_free_context(out_fmt_ctx);

	return 0;
}


